from celery.app.control import Control
import datetime
import threading
import time
from celery.events import EventReceiver
from app import celeryapp

celery_=celeryapp.celery

def import_object(object_path):
    '''
    imports and returns given class string.

    :param object_path: Class path as string
    :type object_path: str

    :returns: Class that has given path
    :rtype: class

    :Example:

    >>> import_object('collections.OrderedDict').__name__
    'OrderedDict'
    '''
    try:
        from importlib import import_module
        module_name='.'.join(object_path.split('.')[:-1])
        mod=import_module(module_name)
        return getattr(mod,object_path.split('.')[-1])
    except Exception as detail:
        raise ImportError(detail)
    
def nested_method(clazz,method,nested):
    from types import CodeType,FunctionType
    '''
    Return the function named <child_name> that is defined inside a <parent> function
    Returns None if nonexistent
    '''
    parent=getattr(clazz,method)
    consts=parent.func_code.co_consts
    for item in consts:
        if isinstance(item,CodeType) and item.co_name==nested:
            return FunctionType(item,globals())
        
class CeleryClient(object):
    _application=None
    _control=None
    _default_queue=None

    def __init__(self):
        self._application=celery_
        self._control=Control(self._application)
        self._default_queue=self._application.amqp.default_queue.name
        self._routes={}

    @property
    def application(self):
        return self._application
    
    @property
    def default_queue(self):
        return self._default_queue
    
    @property
    def routes(self):
        return self._routes
    
    def enable_events(self):
        self._control.enable_events()

    def disable_events(self):
        self._control.disable_events()

    def workers(self):
        response=self._control.inspect().stats()
        if not response:
            return []
        statuses=self.worker_statuses()
        queues=self.active_queues()
        workers=[]
        for name,info in response.items():
            worker=dict()
            worker['name']=name
            worker['status']=statuses[worker['name']]
            worker['concurrency']=info['pool']['max-concurrency']
            worker['broker']={'transport':info['broker']['transport'],'hostname':info['broker']['hostname'],'port':info['broker']['port']}
            worker['queues']=queues[worker['name']]
            workers.append(worker)
        return workers

    def worker_statuses(self):
        '''
        get worker statuses
        :return:
        '''
        response=self._control.ping()
        if not response:
            return []
        workers={}
        for w in response:
            for k,v in w.items():
                for k_inner,v_inner in v.items():
                    if k_inner=='ok' and v_inner=='pong':
                        workers[k]='Active'
                    else:
                        workers[k]='Passive'
                    break
        return workers
    
    def active_queues(self):
        '''
        get queue mappings with workers
        :return:
        '''
        response=self._control.inspect().active_queues()
        if not response:
            return []
        workers={}
        for w,queues in response.items():
            workers[w]=list()
            for q in queues:
                workers[w].append(q['name'])
        return workers
    
    def registered_tasks(self):
        '''
        get registered task list
        :return:
        '''
        response=self._control.inspect().registered()
        if not response:
            return []
        all_tasks=set()
        for worker,tasks in response.items():
            for task in tasks:
                all_tasks.add(task)
        registered_tasks={}
        for task in all_tasks:
            if task in self.routes:
                queue=self.routes[task].get('queue',self.default_queue)
            else:
                queue=self.default_queue
            registered_tasks[task]=queue
        return registered_tasks
    
    def active_tasks(self):
        '''
        get active tasks which is running currently
        :return:
        '''
        response=self._control.inspect().active()
        if not response:
            return []
        tasks=[]
        for worker,task_list in response.items():
            for task in task_list:
                t=dict()
                t['queue']=task['delivery_info']['routing_key']
                t['name']=task['name']
                t['id']=task['id']
                t['worker']=worker
                tasks.append(t)
        return tasks
    
    def reserved_tasks(self):
        '''
        get reserved tasks which is in queue but still waiting to be executed
        :return:
        '''
        response=self._control.inspect().reserved()
        if not response:
            return []
        tasks=[]
        for worker,task_list in response.items():
            for task in task_list:
                t=dict()
                t['queue']=task['delivery_info']['routing_key']
                t['name']=task['name']
                t['id']=task['id']
                t['worker']=worker
                tasks.append(t)
        return tasks
    
    def execute(self,command,parameter):
        def run(*args):
            task_verbose=args[1]
            task=import_object(task_verbose)
            task.delay()

        def revoke(*args):
            ctrl=args[0]
            task_id=args[1]
            ctrl.revoke(task_id,terminate=True,signal='SIGKILL')

        control=self._control
        nested=nested_method(self,'execute',command)
        return nested(*(control,parameter))
    
    def get_manage_bttasks_task_id(self):
        response=self._control.inspect().active()
        if not response:
            return None
        task_id=None
        for worker,task_list in response.items():
            for task in task_list:
                if task['name']=='app.tasks.manage_bttasks':
                    task_id=task['id']
                    break
        return task_id
    
    def purge_celery_tasks(self):
        self._application.connection().channel().queue_purge(self._default_queue)

    def get_bt_exec_task_id(self):
        response=self._control.inspect().active()
        if not response:
            return None
        task_id=None
        for worker,task_list in response.items():
            for task in task_list:
                if task['name']=='app.tasks.exec_bt':
                    task_id=task['id']
                    break
        return task_id
    
class ContextManager(object):
    _client=None
    # _dashboard and _tasks are mutable and same object for each instance
    # so one instance will be used on the scope always
    _dashboard={}
    _events={}
    # todo find a better way to define this field
    EVENT_SIZE_THRESHOLD=100

    def __init__(self,client=None):
        self._client=client or CeleryClient()
    
    @property
    def dashboard(self):
        self._dashboard['workers']=self.workers()
        self._dashboard['registered_tasks']=self.registered_tasks()
        self._dashboard['queue_tasks']=self.queue_tasks()
        return self._dashboard
    
    @property
    def tasks(self):
        _tasks=[]
        for uuid,event in self._events.items():
            task={'uuid':uuid,'name':event.get('name',''),'state':event['type'].replace('task-','').upper(),'args':event.get('args',''),'kwargs':event.get('kwargs',''),'received':event['local_received']}
            _tasks.append(task)
        return _tasks
    
    def add_event(self,event):
        if len(self._events)==self.EVENT_SIZE_THRESHOLD:
            self._events.pop(0)
        if event['uuid'] in self._events:
            exists=self._events[event['uuid']]
            event.update(name=exists.get('name',''))
            event.update(args=exists.get('args',''))
            event.update(kwargs=exists.get('kwargs',''))
        event['local_received']=datetime.datetime.fromtimestamp(event['local_received'])
        self._events.update({event['uuid']:event})
    
    def workers(self):
        return self._client.workers()
    
    def registered_tasks(self):
        return self._client.registered_tasks()

    def queue_tasks(self):
        return {'active':self._client.active_tasks(),'reserved':self._client.reserved_tasks()}
    
class TaskStatus:
    RECEIVED=(0,'received')
    STARTED=(1,'started')

class EventListener(threading.Thread):
    def __init__(self,celery_client,context_manager,enable_events=False):
        threading.Thread.__init__(self)
        self.daemon=True
        self.celery_client=celery_client
        self.context_manager=context_manager
        self.enable_events=enable_events

    def start(self):
        threading.Thread.start(self)

    def run(self):
        ready=False
        while not ready:
            workers=self.celery_client.workers()
            if not workers:
                time.sleep(5)
                continue
            ready=True
            if self.enable_events:
                self.celery_client.enable_events()
            application=self.celery_client.application
            with application.connection() as conn:
                receiver=EventReceiver(conn,handlers={'*':self.on_event},app=application)
                receiver.capture(limit=None,timeout=None,wakeup=True)
    def on_event(self,event):
        if event['type'].startswith('task-'):
            self.context_manager.add_event(event)

celery_client=CeleryClient()
context_manager=ContextManager(client=celery_client)
event_listener=EventListener(celery_client,context_manager,enable_events=True)
event_listener.start()